'''
Author: duliang thinktanker@163.com
Date: 2025-09-11 22:27:13
LastEditors: duliang thinktanker@163.com
LastEditTime: 2025-09-13 21:05:18
FilePath: 
Description: 
'''
from myyolo import MyYOLO
import os
import sys

def recheck_images(model_path, input_folder, output_folder):
    """
    对文件夹中的图片进行重新检测并保存标签文件
    :param model_path: 模型路径
    :param input_folder: 输入图片文件夹路径
    :param output_folder: 输出结果文件夹路径
    """
    # 创建YOLO实例
    yolo = MyYOLO()
    
    # 加载模型
    if os.path.exists(model_path):
        yolo.load_model(model_path)
    else:
        print(f"模型文件 {model_path} 不存在")
        return
    
    # 检查输入文件夹
    if not os.path.exists(input_folder):
        print(f"输入文件夹 {input_folder} 不存在")
        return
    
    # 创建输出文件夹
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    
    # 创建labelimg文件夹用于保存原图和标签文件
    labelimg_folder = os.path.join(output_folder, "labelimg")
    if not os.path.exists(labelimg_folder):
        os.makedirs(labelimg_folder)
    
    # 支持的图片格式
    supported_formats = ('.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif')
    
    # 获取文件夹中的所有图片文件
    image_files = [
        f for f in os.listdir(input_folder)
        if f.lower().endswith(supported_formats)
    ]
    
    if not image_files:
        print(f"文件夹 {input_folder} 中没有找到支持的图片文件")
        return
    
    print(f"找到 {len(image_files)} 个图片文件，开始处理...")
    
    # 处理每个图片文件
    for image_file in image_files:
        try:
            image_path = os.path.join(input_folder, image_file)
            print(f"正在处理: {image_file}")
            
            # 执行检测
            annotated_image, results = yolo.detect_image(image_path)
            
            # 保存结果图像
            output_path = os.path.join(output_folder, f"result_{image_file}")
            yolo.save_result(annotated_image, output_path)
            print(f"结果已保存到: {output_path}")
            # 复制原图到labelimg文件夹
            target_image_path = os.path.join(labelimg_folder, image_file)
            # if os.path.exists(target_image_path):
            #     os.remove(target_image_path)           
            
            # 保存YOLO格式的标签文件
            label_path = os.path.join(labelimg_folder, f"{os.path.splitext(image_file)[0]}.txt")
            # 读取原图获取尺寸信息
            import cv2
            image = cv2.imread(image_path)
            if image is not None:
                yolo.save_yolo_labels(results, label_path, image.shape[1], image.shape[0])
                print(f"标签已保存到: {label_path}")
            # 复制而不是移动原图
            if not os.path.exists(target_image_path):
                import shutil
                shutil.copy(image_path, target_image_path)
                print(f"原图已保存到: {target_image_path}")
            
            
        except Exception as e:
            print(f"处理图片 {image_file} 时出错: {str(e)}")
            continue
    
    print(f"处理完成，共处理 {len(image_files)} 个图片文件")

if __name__ == "__main__":
    # 配置路径
    MODEL_PATH = r"./models/best.pt"  # 模型路径
    INPUT_FOLDER = r"D:\results\labelimg"   # 输入图片文件夹
    OUTPUT_FOLDER = r"D:\results"     # 输出结果文件夹
    
    # 执行检测
    recheck_images(MODEL_PATH, INPUT_FOLDER, OUTPUT_FOLDER)